package org.example.state;

import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;
import org.example.flink.pojo.WaterSensor;
import org.example.function.WaterSensorFunction;

/**
 * @Author: tang
 * @Description:
 * @Date 2025/2/25 11:22
 */
public class OperatorBroadcastStateDemo {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        environment.setParallelism(2);
        environment.setStateBackend(new HashMapStateBackend());

        // 数据流
        SingleOutputStreamOperator<WaterSensor> streamOperator1 = environment
                .socketTextStream("127.0.0.1", 7777)
                .map(new WaterSensorFunction());

        // 广播流
        SingleOutputStreamOperator<String> streamOperator2 = environment
                .socketTextStream("127.0.0.1", 8888);
        MapStateDescriptor<String, Integer> mapStateDescriptor =
                new MapStateDescriptor<>("map-state", Types.STRING, Types.INT);
        BroadcastStream<String> broadcastStream = streamOperator2.broadcast(mapStateDescriptor);

        streamOperator1.connect(broadcastStream)
                .process(new BroadcastProcessFunction<WaterSensor, String, String>() {
                    @Override
                    public void processElement(WaterSensor value, BroadcastProcessFunction<WaterSensor, String, String>.ReadOnlyContext ctx, Collector<String> out) throws Exception {
                        //  只读的
                        ReadOnlyBroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
                        Integer thresholdValue = broadcastState.get("threshold");
                        Integer threshold = (thresholdValue == null ? 0 : thresholdValue);
                        if (value.getVc() > threshold) {
                            out.collect("当前vc: " + value.getVc() + "大于阈值: " + threshold);
                        }
                    }

                    /**
                     * 只有广播流能写
                     */
                    @Override
                    public void processBroadcastElement(String value,
                                                        BroadcastProcessFunction<WaterSensor, String, String>.Context ctx,
                                                        Collector<String> out) throws Exception {
                        // 可写
                        BroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
                        broadcastState.put("threshold",Integer.parseInt(value));
                        out.collect("广播流中阈值更新为: " + value);
                    }
                })
                .print();

        environment.execute();
    }

}
